import os
from mayavi import mlab
import tqdm

from hj_dataset_devkit import HongjingScene

from common_utils import interpolate_obj_anno, show_box, show_np_cloud, Colors, dump_json

def output_scenes_interp(scenes_root: str, output_root: str) -> None:
    def check_scene_name(name: str) -> bool:
        if len(name) >= 19:
            obj_anno_path = os.path.join(scenes_root, name, 'obj_anno')
            if name[:10] > '2023_00_00' and name[:10] < '2023_08_00' and os.path.exists(obj_anno_path):
                return True
        return False

    scene_names = []
    for folder in os.listdir(scenes_root):
        if check_scene_name(folder):
            scene_names.append(folder)
    if not os.path.exists(output_root):
        os.mkdir(output_root)
    for scene_name in tqdm.tqdm(scene_names):
        scene_path = os.path.join(scenes_root, scene_name)
        scene = HongjingScene.from_path(scene_path)
        intep_anno = interpolate_obj_anno(scene)
        if not os.path.exists(os.path.join(output_root, scene_name)):
            os.mkdir(os.path.join(output_root, scene_name))
        anno_folder = os.path.join(output_root, scene_name, 'obj_anno')
        if not os.path.exists(anno_folder):
            os.mkdir(anno_folder)
        for stamp, anno in intep_anno.items():
            anno_path = os.path.join(anno_folder, f'{stamp}.json')
            dump_json(anno_path, [obj.to_dict() for obj in anno])

def interp_check_one_scene(scene_path: str) -> None:
    scene = HongjingScene.from_path(scene_path)
    intep_anno = interpolate_obj_anno(scene)
    for frame in scene.frames[5: -1: 50]:
        cloud = frame.get_merged_lidar_xyzi_cloud()
        figure = mlab.figure(f'{frame.get_timestamp()}', Colors.Dark)
        show_np_cloud(figure, cloud)
        anno = intep_anno[frame.get_timestamp()]
        for obj in anno:
            show_box(figure, obj.box.center.to_list(), obj.box.size.to_list(),
                     obj.box.rotation.to_list(), Colors.Blue)
    mlab.show()

if __name__ == '__main__':
    # interp_check_one_scene('$SCENE_PATH')
    # output_scenes_interp('$SCENES_ROOT', '$OUTPUT_ROOT')
